﻿
namespace Apache.Cassandra.Client {
    using System;
    using System.Collections;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Linq;
    using System.Net;
    using System.Net.Sockets;
    using System.Security;
    using System.Threading;
    using Apache.Cassandra.Client.Diagnostics;

    internal class CassandraConnectionPool {

        private enum State {
            Initializing,
            Running,
            ShuttingDown, 
        }

        internal class CassandraEndPointPool {
            private State state;
            private readonly IPAddress address;
            private readonly Semaphore poolSemaphore;
            private readonly List<CassandraInternalConnection> connections;
            private readonly Queue<CassandraInternalConnection> stack;
            private int count = 0;

            internal CassandraEndPointPool(IPAddress address, int size) {
                this.state = State.Initializing;

                this.address = address;
                this.poolSemaphore = new Semaphore(0, size);
                this.connections = new List<CassandraInternalConnection>(size);
                this.stack = new Queue<CassandraInternalConnection>(size);

                this.state = State.Running;
            }

            internal IPAddress Address {
                get {
                    return address;
                }
            }

            internal int Count {
                get {
                    return count;
                }
            }

            internal Semaphore PoolSemaphore {
                get {
                    return poolSemaphore;
                }
            }

            internal CassandraInternalConnection Pop() {
                lock (stack) {
                    return 0 == stack.Count ? null : stack.Dequeue();
                }
            }

            internal bool Push(CassandraInternalConnection value) {
                if (State.Running != state) {
                    return false;
                }

                lock (stack) {
                    stack.Enqueue(value);
                }

                return true;
            }

            internal void AddConnection(CassandraInternalConnection connection) {
                lock (connections) {
                    connections.Add(connection);
                    count = connections.Count;
                }
            }

            internal bool RemoveConnection(CassandraInternalConnection connection) {
                lock (connections) {
                    bool removed = connections.Remove(connection);
                    count = connections.Count;
                    return removed;
                }
            }

            internal List<CassandraInternalConnection> ReclaimEmancipatedObjects() {
                List<CassandraInternalConnection> reclaimedObjects = new List<CassandraInternalConnection>();

                lock (connections) {
                    int count = connections.Count;

                    for (int i = 0; i < count; ++i) {
                        CassandraInternalConnection connection = connections[i];
                        if (null == connection) {
                            continue;
                        }

                        bool locked = false;
                        try {
                            Monitor.TryEnter(connection, ref locked);

                            if (locked) {
                                if (connection.IsEmancipated) {
                                    connection.PrePush(null);
                                    reclaimedObjects.Add(connection);
                                }
                            }
                        } finally {
                            if (locked) {
                                Monitor.Exit(connection);
                            }
                        }
                    }
                }

                return reclaimedObjects;
            }

            internal void DoomConnections() {
                lock(connections) {
                    int count = connections.Count;

                    for (int i = 0; i < count; ++i) {
                        CassandraInternalConnection connection = connections[i];

                        if (null != connection) {
                            connection.DoomConnection();
                        }
                    }
                }
            }

            internal void Shutdown() {
                state = State.ShuttingDown;
            }
        }

        private interface IEndPointSelector {
            int Next();
        }

        private class RoundRobinEndPointSelector : IEndPointSelector {
            private int current;

            internal RoundRobinEndPointSelector()
                : this(0) {
            }

            internal RoundRobinEndPointSelector(int seed) {
                current = seed - 1;
            }

            public int Next() {
                return Interlocked.Increment(ref current);
            }
        }

        private class RandomEndPointSelector : IEndPointSelector {
            private readonly Random rng = new Random();

            public int Next() {
                lock (rng) {
                    return rng.Next();
                }
            }
        }

        private class IPAddressComparer : IComparer<IPAddress> {
            public int Compare(IPAddress x, IPAddress y) {
                if (x == y) {
                    return 0;
                }
                if (null == x) {
                    return -1;
                }
                if (null == y) {
                    return 1;
                }

                Debug.Assert(AddressFamily.InterNetwork == x.AddressFamily);
                Debug.Assert(AddressFamily.InterNetwork == y.AddressFamily);

                int c = 0;

                c = ((IStructuralComparable)x.GetAddressBytes()).CompareTo(y.GetAddressBytes(), Comparer.Default);
                if (0 != c) {
                    return c;
                }

                return 0;
            }
        }

        private State state;
        private readonly Random rng = new Random();

        private readonly CassandraConnectionString settings;
        private readonly CassandraConnectionFactory connectionFactory;
        private readonly CassandraClientPoolCounters performanceCounters;

        // Cluster connections
        private readonly ReaderWriterLockSlim endpointLock = new ReaderWriterLockSlim();
        private readonly IEndPointSelector endpointSelector;
        private int endpointCount = 0;
        private readonly IDictionary<IPAddress, CassandraEndPointPool> endpointPools = new SortedDictionary<IPAddress, CassandraEndPointPool>(new IPAddressComparer());
        private readonly IList<CassandraEndPointPool> inactiveEndPointPools = new List<CassandraEndPointPool>();
        private int managedConnections = 0;

        private readonly WaitCallback createRequest;

        // Wait handles
        private int waitCount = 0;
        private readonly Semaphore creationSemaphore = new Semaphore(1, 1);

        // Deactivate queue
        private readonly Queue<CassandraInternalConnection> deactivateQueue;
        private readonly WaitCallback deactivateCallback; 

        // Clean up timer
        private readonly TimeSpan cleanupWait;
        private Timer cleanupTimer;

        // End point pruning timer
        private readonly TimeSpan pruningWait;
        private Timer pruningTimer;

        internal CassandraConnectionPool(CassandraConnectionString settings) {
            state = State.Initializing;

            this.settings = settings;

            this.performanceCounters = CreatePerformanceCounters();

            foreach (IPAddress address in ServerAddresses) {
                endpointPools.Add(address, new CassandraEndPointPool(address, MaxPoolSize));
                PerformanceCounters.NumberOfActiveEndPointPools.Increment();
            }
            endpointCount = endpointPools.Count;

            cleanupWait = TimeSpan.FromSeconds(120);
            pruningWait = TimeSpan.FromSeconds(30);

            endpointSelector = CreateEndpointSelector();
            connectionFactory = CassandraConnectionFactory.Factory;

            if (UseDeactivateQueue) {
                deactivateQueue = new Queue<CassandraInternalConnection>();
                deactivateCallback = new WaitCallback(ProcessDeactivateQueue);
            }

            createRequest = new WaitCallback(CreateRequest);

            state = State.Running;
        }

        private CassandraClientPoolCounters PerformanceCounters {
            get {
                return performanceCounters;
            }
        }

        private ISet<IPAddress> ServerAddresses {
            get {
                ISet<IPAddress> addresses = new HashSet<IPAddress>();

                foreach (string server in settings.Server.Split((char[])null, StringSplitOptions.RemoveEmptyEntries)) {
                    foreach (IPAddress address in Dns.GetHostAddresses(server)) {
                        addresses.Add(address);
                    }
                }

                return addresses;
            }
        }

        private TimeSpan ConnectTimeout {
            get {
                // TODO: Return configuration value
                return TimeSpan.FromMilliseconds(3000);
            }
        }
        private TimeSpan CreationTimeout {
            get {
                // TODO: Return configuration value
                return TimeSpan.FromMilliseconds(1000);
            }
        }

        private int MaxPoolSize {
            get {
                // TODO: Return configuration value
                return 128;
            }
        }

        private int MinPoolSize {
            get {
                // TODO: Return configuration value
                return 0;
            }
        }
 
        private bool UseDeactivateQueue {
            get {
                return settings.UseDeactivateQueue;
            }
        } 

        internal void Startup() { 
            cleanupTimer = CreateCleanupTimer(); 
            pruningTimer = CreatePruningTimer();

            BringToSize();
        }

        private CassandraClientPoolCounters CreatePerformanceCounters() {
            string instance;

            if (String.IsNullOrEmpty(settings.PerformanceCountersInstance)) {
                instance = CassandraConnectionFactory.PerformanceCountersInstanceName;
            } else {
                instance = String.Format("{0}[{1}]", CassandraConnectionFactory.PerformanceCountersInstanceName, settings.PerformanceCountersInstance);
            }

            return new CassandraClientPoolCounters(instance);
        }

        private IEndPointSelector CreateEndpointSelector() {
            // TODO: Initialize endpoint selector from configuration file
            lock (rng) {
                return new RoundRobinEndPointSelector(rng.Next());
            }
        }

        private Timer CreateCleanupTimer() {
            return (new Timer(new TimerCallback(this.CleanupCallback), null, cleanupWait, cleanupWait)); 
        }

        private Timer CreatePruningTimer() {
            return (new Timer(new TimerCallback(this.PruningCallback), null, pruningWait, pruningWait)); 
        }

        private Semaphore CreationSemaphore {
            get {
                return creationSemaphore;
            }
        }
 
        private void CleanupCallback(Object state) {
            // Destroy connections over minimum pool size
            endpointLock.EnterReadLock();
            try {
                foreach (var i in endpointPools) {
                    CassandraEndPointPool pool = i.Value;

                    // TODO: Do not destroy connections under minimum TTL
                    while (pool.Count > MinPoolSize) {
                        if (!pool.PoolSemaphore.WaitOne(0, false)) {
                            break;
                        }
                        CassandraInternalConnection connection = pool.Pop();
                        Debug.Assert(null != connection);
                        PerformanceCounters.NumberOfFreeConnections.Decrement();
                        DestroyConnection(pool, connection);
                    }
                }
            } finally {
                endpointLock.ExitReadLock();
            }

            // Queue up a request to bring us up to MinPoolSize 
            BringToSize();
        }

        private void PruningCallback(Object state) {
            // Clean up inactive end point pools
            CassandraEndPointPool[] pools;
            lock (inactiveEndPointPools) {
                pools = new CassandraEndPointPool[inactiveEndPointPools.Count];
                inactiveEndPointPools.CopyTo(pools, 0);
            }
            foreach (CassandraEndPointPool pool in pools) {
                ClearEndPointPool(pool);

                if (0 == pool.Count) {
                    bool removed;
                    lock (inactiveEndPointPools) {
                        removed = inactiveEndPointPools.Remove(pool);
                    }
                    if (removed) {
                        PerformanceCounters.NumberOfInactiveEndPointPools.Decrement();
                    }
                }
            }

            // Check offline end points
            ISet<IPAddress> addresses = ServerAddresses;
            endpointLock.EnterReadLock();
            try {
                addresses.ExceptWith(endpointPools.Keys);
            } finally {
                endpointLock.ExitReadLock();
            }

            foreach (IPAddress address in addresses) {
                CassandraInternalConnection connection = null;

                try {
                    IPEndPoint endpoint = new IPEndPoint(address, settings.Port);
                    connection = connectionFactory.CreateConnection(null, endpoint, settings);
                    if (null == connection) {
                        break;
                    }

                    CassandraEndPointPool pool = null;

                    endpointLock.EnterUpgradeableReadLock();
                    try {
                        if (!endpointPools.ContainsKey(address)) {
                            pool = new CassandraEndPointPool(address, MaxPoolSize);
                            endpointLock.EnterWriteLock();
                            try {
                                endpointPools.Add(address, pool);
                            } finally {
                                endpointLock.ExitWriteLock();
                            }
                            PerformanceCounters.NumberOfActiveEndPointPools.Increment();
                        }
                    } finally {
                        endpointLock.ExitUpgradeableReadLock();
                    }

                    if (null != pool) {
                        QueuePoolCreateRequest(pool);
                    }
                } catch (Exception e) {
                    if (!e.IsCatchableType()) {
                        throw;
                    }
                } finally {
                    if (null != connection) {
                        connection.Dispose();
                    }
                }
            }
        }

        private void BringToSize() {
            endpointLock.EnterReadLock();
            try {
                foreach (var i in endpointPools) {
                    CassandraEndPointPool pool = i.Value;

                    int count = pool.Count;
                    if (count >= MaxPoolSize) {
                        continue;
                    }
                    if (count < MinPoolSize) {
                        QueuePoolCreateRequest(i.Value);
                    }
                }
            } finally {
                endpointLock.ExitReadLock();
            }
        }

        private void DeactivateConnection(CassandraInternalConnection connection) {
            connection.Deactivate();

            if (State.ShuttingDown == state || connection.IsDoomed) {
                IPEndPoint endpoint = connection.EndPoint as IPEndPoint;
                CassandraEndPointPool pool = null;

                if (null != endpoint) {
                    endpointLock.EnterReadLock();
                    try {
                        endpointPools.TryGetValue(endpoint.Address, out pool);
                    } finally {
                        endpointLock.ExitReadLock();
                    }
                }
                DestroyConnection(pool, connection);

                if (null != pool) {
                    QueuePoolCreateRequest(pool);
                }
            } else {
                Debug.Assert (State.Running == state);
                PushNewConnection(connection);
            }
        }

        internal void DestroyConnection(CassandraInternalConnection connection) {
            IPEndPoint endpoint = connection.EndPoint as IPEndPoint;
            CassandraEndPointPool pool = null;

            if (null != endpoint) {
                endpointLock.EnterReadLock();
                try {
                    endpointPools.TryGetValue(endpoint.Address, out pool);
                } finally {
                    endpointLock.ExitReadLock();
                }
            }

            DestroyConnection(pool, connection);
        }

        private void DestroyConnection(CassandraEndPointPool pool, CassandraInternalConnection connection) {
            bool removed = false;

            if (null != pool) {
                IPEndPoint endpoint = connection.EndPoint as IPEndPoint;
                Debug.Assert(pool.Address == endpoint.Address && settings.Port == endpoint.Port);
                removed = pool.RemoveConnection(connection);
                Debug.Assert(removed, "attempt to destroy connection not in list");
            }

            if (removed) {
                Interlocked.Decrement(ref managedConnections);
                PerformanceCounters.NumberOfPooledConnections.Decrement();
            }
            connection.Dispose();

            PerformanceCounters.HardDisconnectsPerSecond.Increment(); 
        }

        internal void PushNewConnection(CassandraInternalConnection connection) {
            Debug.Assert(null != connection, "Adding null connection to the pool.");
            Debug.Assert(connection.CanBePooled, "Adding connection that cannot be pooled to the pool.");

            IPEndPoint endpoint = connection.EndPoint as IPEndPoint;
            Debug.Assert(null != endpoint, "Adding connection with an incompatible endpoint to the pool.");

            CassandraEndPointPool pool;
            endpointLock.EnterReadLock();
            try {
                endpointPools.TryGetValue(endpoint.Address, out pool);
            } finally {
                endpointLock.ExitReadLock();
            }

            if (null != pool && pool.Push(connection)) {
                pool.PoolSemaphore.Release(1);
                PerformanceCounters.NumberOfFreeConnections.Increment();
            } else {
                DestroyConnection(pool, connection);
            }
        }

        public void PushConnection(CassandraInternalConnection connection, CassandraConnection owner) {
            Debug.Assert(null != connection, "Pushing null connection to the pool.");

            PerformanceCounters.SoftDisconnectsPerSecond.Increment();

            lock (connection) {
                connection.PrePush(owner);
            }

            if (UseDeactivateQueue) {
                PerformanceCounters.NumberOfStasisConnections.Increment();
                bool needToQueueWorkItem;

                lock (deactivateQueue) {
                    needToQueueWorkItem = 0 == deactivateQueue.Count;
                    deactivateQueue.Enqueue(connection);
                }

                if (needToQueueWorkItem) {
                    ThreadPool.QueueUserWorkItem(deactivateCallback);
                }
            } else {
                DeactivateConnection(connection);
            }
        }

        public CassandraInternalConnection PopConnection(CassandraConnection owner) {
            PerformanceCounters.SoftConnectsPerSecond.Increment();

            if(state != State.Running) { 
                return null; 
            }

            if (endpointCount == 0) {
                return null;
            }

            int preferred = endpointSelector.Next();

            int count;
            IPAddress[] keys;
            WaitHandle[] handles;

            endpointLock.EnterReadLock();
            try {
                count = endpointPools.Count;
                keys = new IPAddress[count];
                handles = new WaitHandle[count + 1];

                int index = count - preferred % count;
                foreach (var i in endpointPools) {
                    keys[index % count] = i.Key;
                    handles[index % count] = i.Value.PoolSemaphore;
                    ++index;
                }
            } finally {
                endpointLock.ExitReadLock();
            }
            handles[count] = CreationSemaphore;

            Interlocked.Increment(ref waitCount);

            CassandraInternalConnection connection = null;
            Stopwatch stopwatch = Stopwatch.StartNew();
            TimeSpan timeout = ConnectTimeout;
            do {
                int handle = WaitHandle.WaitAny(handles, timeout, false);

                if (WaitHandle.WaitTimeout == handle) {
                    Interlocked.Decrement(ref waitCount);
                    return null;
                } else if (handle == count) {
                    // Create new connection
                    try {
                        CassandraEndPointPool pool = EndPointPoolAt(preferred);
                        if (null != pool) {
                            connection = UserCreateRequest(pool, owner);
                        }
                    } catch {
                        if (null == connection) {
                            Interlocked.Decrement(ref waitCount);
                        }
                        throw;
                    } finally {
                        CreationSemaphore.Release();

                        if (null != connection) {
                            Interlocked.Decrement(ref waitCount);
                        }
                    }

                    if (null == connection) {
                        // Do not try to create more connections when pool full
                        if (managedConnections >= MaxPoolSize * endpointCount && 0 != MaxPoolSize) {
                            if (!ReclaimEmancipatedObjects()) {
                                Array.Resize(ref handles, count);
                            }
                        }
                    }
                } else {
                    // Connection available
                    CassandraEndPointPool pool = null;
                    endpointLock.EnterReadLock();
                    try {
                        if (endpointPools.TryGetValue(keys[handle], out pool)) {
                            Interlocked.Decrement(ref waitCount);
                        }
                    } finally {
                        endpointLock.ExitReadLock();
                    }
                    if (null != pool) {
                        connection = pool.Pop();
                        if (null != connection) {
                            PerformanceCounters.NumberOfFreeConnections.Decrement();
                        }
                    }
                }

                // Update wait timeout
                if (null == connection) {
                    timeout = ConnectTimeout - stopwatch.Elapsed;
                    if (timeout < TimeSpan.Zero) {
                        Interlocked.Decrement(ref waitCount);
                        return null;
                    }
                }
            } while (null == connection);

            if (null != connection) {
                lock (connection) {
                    connection.PostPop(owner);
                }
                try {
                    connection.Activate();
                } catch (SecurityException) {
                    this.PushConnection(connection, owner);
                    throw;
                }
            }

            return connection;
        }

        private CassandraInternalConnection CreateConnection(CassandraEndPointPool pool, CassandraConnection owner) {
            Debug.Assert(null != pool);

            CassandraInternalConnection connection = null;
            do {
                try {
                    IPEndPoint endpoint = new IPEndPoint(pool.Address, settings.Port);
                    connection = connectionFactory.CreateConnection(this, endpoint, settings);
                    if (null == connection) {
                        throw new InvalidOperationException();
                    }
                    PerformanceCounters.HardConnectsPerSecond.Increment();
                    connection.MakePooledConnection(this);

                    if (!connection.CanBePooled) { 
                        throw new InvalidOperationException();
                    }

                    connection.PrePush(null);

                    pool.AddConnection(connection);
                    Interlocked.Increment(ref managedConnections);
                    PerformanceCounters.NumberOfPooledConnections.Increment();

                    return connection;
                } catch (Exception e) {
                    if (!e.IsCatchableType()) {
                        throw;
                    }

                    OfflineEndpointPool(pool);
                }

                pool = EndPointPoolAt(endpointSelector.Next());
            } while (null == connection && pool != null);

            return null;
        }

        private CassandraEndPointPool EndPointPoolAt(int index) {
            endpointLock.EnterReadLock();
            try {
                int count = endpointPools.Count;
                if (0 == count) {
                    return null;
                }
                for (int i = 0; i < count; ++i) {
                    CassandraEndPointPool pool = endpointPools.ElementAt((i + index) % count).Value;
                    if (pool.Count < MaxPoolSize) {
                        return pool;
                    }
                }
            } finally {
                endpointLock.ExitReadLock();
            }

            return null;
        }

        private void ProcessDeactivateQueue(object state) {
            CassandraInternalConnection[] queue;

            lock (deactivateQueue) {
                queue = deactivateQueue.ToArray();
                deactivateQueue.Clear();
            }

            foreach (CassandraInternalConnection connection in queue) {
                PerformanceCounters.NumberOfStasisConnections.Decrement();
                DeactivateConnection(connection);
            }
        }

        private void QueuePoolCreateRequest(CassandraEndPointPool pool) {
            Debug.Assert(null != pool);

            if (State.Running == state) {
                ThreadPool.QueueUserWorkItem(createRequest, pool);
            }
        }

        private CassandraInternalConnection UserCreateRequest(CassandraEndPointPool pool, CassandraConnection owner) {
            if (pool.Count < MaxPoolSize || 0 == MaxPoolSize) {
                return CreateConnection(pool, owner);
            }
            return null;
        }

        private void CreateRequest(object state) {
            if (State.Running != this.state) {
                return;
            }

            CassandraEndPointPool pool = state as CassandraEndPointPool;
            Debug.Assert(null != pool);

            pool.ReclaimEmancipatedObjects();

            if (!(pool.Count < MinPoolSize)) {
                return;
            }

            if (CreationSemaphore.WaitOne(CreationTimeout, false)) {
                try {
                    while (pool.Count < MinPoolSize) {
                        CassandraInternalConnection connection = CreateConnection(pool, (CassandraConnection)null);

                        if (null == connection) {
                            break;
                        }
                        PushNewConnection(connection);
                    }
                } finally {
                    CreationSemaphore.Release();
                }
            } else {
                QueuePoolCreateRequest(pool);
            }
        }

        private bool ReclaimEmancipatedObjects() {
            List<CassandraInternalConnection> reclaimedObjects = new List<CassandraInternalConnection>();

            endpointLock.EnterReadLock();
            try {
                foreach (var p in endpointPools) {
                    reclaimedObjects.AddRange(p.Value.ReclaimEmancipatedObjects());
                }
            } finally {
                endpointLock.ExitReadLock();
            }

            int count = reclaimedObjects.Count;
            for (int i = 0; i < count; ++i) {
                PerformanceCounters.NumberOfReclaimedConnections.Increment();
                DeactivateConnection(reclaimedObjects[i]);
            }

            return count > 0;
        }

        private void OfflineEndpointPool(CassandraEndPointPool pool) {
            bool removed;

            endpointLock.EnterWriteLock();
            try {
                removed = endpointPools.Remove(pool.Address);
            } finally {
                endpointLock.ExitWriteLock();
            }

            if (!removed) {
                return;
            }

            PerformanceCounters.NumberOfActiveEndPointPools.Decrement();

            pool.Shutdown();
            pool.DoomConnections();

            Interlocked.Decrement(ref endpointCount);
            Interlocked.Add(ref managedConnections, 0 - pool.Count);

            lock (inactiveEndPointPools) {
                inactiveEndPointPools.Add(pool);
            }

            PerformanceCounters.NumberOfInactiveEndPointPools.Increment();
        }

        private void ClearEndPointPool(CassandraEndPointPool pool) {
            CassandraInternalConnection connection;

            // Clean up free connections
            while (null != (connection = pool.Pop())) {
                PerformanceCounters.NumberOfFreeConnections.Decrement();
                pool.RemoveConnection(connection);
                connection.Dispose();
            }

            // Finally, reclaim emancipated connections
            List<CassandraInternalConnection> reclaimedObjects = pool.ReclaimEmancipatedObjects();
            int count = reclaimedObjects.Count;
            for (int i = 0; i < count; ++i) {
                PerformanceCounters.NumberOfReclaimedConnections.Increment();
                connection = reclaimedObjects[i];

                pool.RemoveConnection(connection);
                connection.Dispose();
            }
        }
    }
}
